from backtrader import indicator, position
import pandas as pd
from pymongo import MongoClient
from datetime import datetime
import time
import math
import talib
import backtrader as bt
import backtrader.indicators as btind
import backtrader.feeds as btfeeds

import sys
sys.path.append('../')
from target.indicator_ma import *
from my_plot import *

objJson = {
    "_id": 'uuidstr',
    "adjustflag": 'adjustflag',
    "code": 'code',
    'start': 'start',
    'end': 'end',
    'startcash': 'startcash',
    'endtcash': 'cerebro.broker.getvalue()',
    'com': 'com',
    'qts': 'qts',
    'basicsData': [],  # 基础数据
    'isbuyArr': [],  # 买入一笔 买入价格
    'issellArr': [],  # 卖出一笔  卖出价格
    'pnlcommArr': [],  # 一笔交易后的盈亏
    'currentCapital': [],  # 资金走势
    'indicatorLines': {},  # 用户所选指标线

    'annualReturn': [],  # 年度回报
    'Calmar': [],  # 卡码率
    'DrawDown': {},  # 回撤
    'sharperatio': math.nan,  # 夏普比率
    'sqnIndex': {},  # SQN指数
    'TimeReturn': [],  # 基于时间范围内回报
}



# 创建策略继承bt.strategy
class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
        # 判断是否输出该日志
        ('printlog', False),
    )

    # 初始化部分指标

    def __init__(self):

        # 指标进行注册
        self.indicators = {}
        # 初始化注册指标
        for param in indicator_params:
            if param['code'] == 'ma':
                indicator = bt.indicators.MovingAverageSimple(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'ema':
                indicator = bt.indicators.ExponentialMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'smma':
                indicator = bt.indicators.SmoothedMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'wma':
                indicator = bt.indicators.WeightedMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'dema':
                indicator = bt.indicators.DoubleExponentialMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'kama':
                indicator = bt.indicators.AdaptiveMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'zlema':
                indicator = bt.indicators.ZeroLagExponentialMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'hma':
                indicator = bt.indicators.HullMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'zlind':
                indicator = bt.indicators.ZeroLagIndicator(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'dma':
                indicator = bt.indicators.DicksonMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator
            
            elif param['code'] == 'highest':
                indicator = bt.indicators.Highest(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator
            
            elif param['code'] == 'lowest':
                indicator = bt.indicators.Lowest(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'bbandsTop':
                indicator = bt.ind.BBands(
                    param['data'], period=param['params'][0],devfactor=2.0).top
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator
            
            elif param['code'] == 'bbandsMid':
                indicator = bt.ind.BBands(
                    param['data'], period=param['params'][0],devfactor=2.0).mid
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator
            
            elif param['code'] == 'bbandsBot':
                indicator = bt.ind.BBands(
                    param['data'], period=param['params'][0],devfactor=2.0).bot
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator
        # 初始化指标lines对象(用于存储绘图数据)
        for item in indicator_user_arr:
            objJson['indicatorLines'][item] = []

        # 跟踪挂单
        self.order = None
        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None

        self.h = self.datas[0].high
        self.l = self.datas[0].low
        self.o = self.datas[0].open
        self.c = self.datas[0].close

        
        self.bbTop = bt.ind.BBands(self.data, period=10,devfactor=2.0).top
        self.bbMid = bt.ind.BBands(self.data, period=10,devfactor=2.0).mid
        self.bbBot = bt.ind.BBands(self.data, period=10,devfactor=2.0).bot

        self.real = bt.indicators.ParabolicSAR(self.data,period=4,af=0.02,afmax=0.2)

        self.maxh = bt.indicators.Highest(self.datas[0].high, period=5)
        self.maxl = bt.indicators.Highest(self.datas[0].close, period=5)
        #买卖状态(1:做多；-1:做空；0:不交易;)
        self.Dealstatus = 0

    # 策略核心代码

    def next(self):

        current_date = self.datas[0].datetime.date(0).isoformat()  # 当前日期
        current_capital = self.broker.getvalue()  # 当前剩余资金
        print('当前日期:{},close:{},SAR:{}'.format(current_date,self.c[0],self.real[0]))

        if(self.maxh[0] >= self.h[-1]):
            self.Dealstatus += 1
        if(self.maxl[0] <= self.l[-1]):
            self.Dealstatus -=1
        
        if(self.c[0] <= self.bbBot[0]):
            self.Dealstatus +=1
        if(self.c[0] >= self.bbTop[0]):
            self.Dealstatus -=1
        
        if(self.real[0] > self.c[0] and self.real[-1] < self.c[-1]):
            self.Dealstatus +=1
        if(self.real[0] < self.c[0] and self.real[-1] > self.c[-1]):
            self.Dealstatus -=1
        

        if(self.Dealstatus > 0):
            self.buy(size=1)
        elif(self.Dealstatus == 0):
            pass
        else:
            self.sell(size=1)
        #     # 判断是否有持仓
        #     if self.position:
        #         self.sell(size=1)
        #     else:
        #         self.sell(size=1)
        
        print(self.Dealstatus)
        self.Dealstatus = 0
        

    # 日志输出

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            # 记录策略的执行日志
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')

    # 订单状态通知，买入卖出都是下单
    def notify_order(self, order):
        # 提交了/接受了,  买/卖订单什么都不做
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 检查一个订单是否完成
        # 注意:当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                print('买入价格:{}---买入手续费{}'.format(order.executed.price,
                      order.executed.comm))
                objJson['isbuyArr'].append(
                    {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            elif order.issell():
                print('卖出价格:{}---卖出手续费{}'.format(order.executed.price,
                      order.executed.comm))
                objJson['issellArr'].append(
                    {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消/保证金不足/拒绝', doprint=False)
        # 其他状态记录为：无法挂单
        self.order = None

    # 交易状态通知，一买一卖算交易（交易净利润）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        objJson['pnlcommArr'].append(
            {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(trade.pnlcomm, 2)})

    # 策略结束时，多用于参数调优
    def stop(self):
        # 回测结束存储数据
        # print(objJson)
        self.log('(MA均线： %2d日) 期末总资金 %.2f' %
                 (self.params.maperiod, self.broker.getvalue()), doprint=False)


def main(execution_id, adjustflag, code, start, end='', startcash=10000000, com=0.0003, qts=100):

    client = MongoClient('47.100.19.231', 27018)
    db = client.admin
    db.authenticate("quant", "Qweasd123")
    db = client.quant
    mycol = db.t_strategy_execution_result
    mydict = {'_id': execution_id}
    s = mycol.find(mydict).count()
    if s > 0:
        return execution_id

    # 创建主控制器
    cerebro = bt.Cerebro(stdstats= True)
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)
    # cerebro.optstrategy(TestStrategy,maperiod=range(10, 15))
    # 获取数据
    # df = selectSotck(adjustflag, code, start, end)
    df = writeExcl()
    df.index = pd.to_datetime(df.date)
    df = df[['open', 'high', 'low', 'close', 'volume']]
    # 将数据加载至回测系统
    data = bt.feeds.PandasData(dataname=df)

    cerebro.adddata(data)
    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # cerebro.broker.setcommission(PerOrder(buy_cost=0.0003, sell_cost=0.0013, min_cost=5))
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    cerebro.broker.set_coc(True)

    print('期初总资金: %.2f' % cerebro.broker.getvalue())

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name= '_TimeReturn')
    cerebro.addobserver(my_BuySell)
    result = cerebro.run()
    print('期末总资金: %.2f' % cerebro.broker.getvalue())

    # # 提取收益序列
    # pnl = pd.Series(result[0].analyzers._TimeReturn.get_analysis())
    # # 计算累计收益
    # cumulative = (pnl + 1).cumprod()
    # # 计算回撤序列
    # max_return = cumulative.cummax()
    # drawdown = (cumulative - max_return) / max_return
    # # 计算收益评价指标
    # import pyfolio as pf
    # # 按年统计收益指标
    # perf_stats_year = (pnl).groupby(pnl.index.to_period('y')).apply(lambda data: pf.timeseries.perf_stats(data)).unstack()
    # # 统计所有时间段的收益指标
    # perf_stats_all = pf.timeseries.perf_stats((pnl)).to_frame(name='all')
    # perf_stats = pd.concat([perf_stats_year, perf_stats_all.T], axis=0)
    # perf_stats_ = round(perf_stats,4).reset_index()
    # # 绘制图形
    # import matplotlib.pyplot as plt
    # plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    # import matplotlib.ticker as ticker # 导入设置坐标轴的模块
    # plt.style.use('dark_background') # plt.style.use('dark_background')
    # fig, (ax0, ax1) = plt.subplots(2,1, gridspec_kw = {'height_ratios':[1.5, 4]}, figsize=(20,8))
    # cols_names = ['date', 'Annual\nreturn', 'Cumulative\nreturns', 'Annual\nvolatility',
    #     'Sharpe\nratio', 'Calmar\nratio', 'Stability', 'Max\ndrawdown',
    #     'Omega\nratio', 'Sortino\nratio', 'Skew', 'Kurtosis', 'Tail\nratio',
    #     'Daily value\nat risk']
    # # 绘制表格
    # ax0.set_axis_off() # 除去坐标轴
    # table = ax0.table(cellText = perf_stats_.values,
    #                 bbox=(0,0,1,1), # 设置表格位置， (x0, y0, width, height)
    #                 rowLoc = 'right', # 行标题居中
    #                 cellLoc='right' ,
    #                 colLabels = cols_names, # 设置列标题
    #                 colLoc = 'right', # 列标题居中
    #                 edges = 'open' # 不显示表格边框
    #                 )
    # table.set_fontsize(13)

    # # 绘制累计收益曲线
    # ax2 = ax1.twinx()
    # ax1.yaxis.set_ticks_position('right') # 将回撤曲线的 y 轴移至右侧
    # ax2.yaxis.set_ticks_position('left') # 将累计收益曲线的 y 轴移至左侧
    # # 绘制回撤曲线
    # drawdown.plot.area(ax=ax1, label='drawdown (right)', rot=0, alpha=0.3, fontsize=13, grid=False)
    # # 绘制累计收益曲线
    # (cumulative).plot(ax=ax2, color='#F1C40F' , lw=3.0, label='cumret (left)', rot=0, fontsize=13, grid=False)
    # # 不然 x 轴留有空白
    # ax2.set_xbound(lower=cumulative.index.min(), upper=cumulative.index.max())
    # # 主轴定位器：每 5 个月显示一个日期：根据具体天数来做排版
    # ax2.xaxis.set_major_locator(ticker.MultipleLocator(100))
    # # 同时绘制双轴的图例
    # h1,l1 = ax1.get_legend_handles_labels()
    # h2,l2 = ax2.get_legend_handles_labels()
    # plt.legend(h1+h2,l1+l2, fontsize=12, loc='upper left', ncol=1)

    # fig.tight_layout() # 规整排版
    # plt.show()

    cerebro.plot(
        # plotter= None, # 包含各种绘图属性的对象或类，如果为None，默认取 PlotScheme 类
        numfigs= 1, # 是否将图形拆分成多幅图展示，如果时间区间比较长，建议分多幅展示
        style= 'candel', # 设置主图行情数据的样式为蜡烛图
        # lcolors=mycolors  ,
        plotdist= 0.1, # 设置图形之间的间距
        barup = 'red', bardown= '#45b97c', # 设置蜡烛图上涨和下跌的颜色
        volup= 'red', voldown= '#45b97c', # 设置成交量在行情上涨和下跌情况下的颜色
    )
    
    

    objJson['_id'] = execution_id
    objJson['adjustflag'] = adjustflag
    objJson['code'] = code
    objJson['start'] = start
    objJson['end'] = end
    objJson['startcash'] = startcash
    objJson['endtcash'] = cerebro.broker.getvalue()
    objJson['com'] = com
    objJson['qts'] = qts
    # mycol.insert_one(objJson)

    return execution_id


# 时间处理
def loadDate(str):
    date = datetime.strptime(str, '%Y%m%d').strftime('%Y-%m-%d')
    return (date)


def datestr2num(s):
    return datetime.strptime(s.decode('ascii'), "%d-%m-%Y").date().weekday()



def writeExcl():
    df = pd.read_excel(r'D:\AGL8test.xls')
    df['code'] = 'AGL8'
    return df


# 数据查询
def selectSotck(adjustflag, code, startDate, endDate):
    client = MongoClient('47.100.19.231', 27018)
    db = client.admin
    db.authenticate("quant", "Qweasd123")
    db = client.quant
    mycol = db.testKline

    myquery = {
        "adjustflag": adjustflag,  # 复权类型
        "code": code,  # 股票代码
        "date": {'$gt': startDate, '$lt': endDate}  # 时间区间
    }

    mydoc = mycol.find(myquery)
    data_list = []
    for x in mydoc:
        data_list.append(x)
    objJson['basicsData'] = data_list
    fields = ['date', 'open', 'close', 'high', 'low', 'volume', 'code']
    df = pd.DataFrame(data_list, columns=fields)
    df['date'] = df['date'].map(loadDate)
    test_list = ['open', 'close', 'high', 'low', 'volume']
    df[test_list] = df[test_list].astype('float64')
    # print(df)
    return df


main('1','1','sz.002998', '20201008', '20210630')





class PlotScheme(object):
    def __init__(self):
    # to have a tight packing on the chart wether only the x axis or also
    # the y axis have (see matplotlib)
        self.ytight = False
        # y-margin (top/bottom) for the subcharts. This will not overrule the
        # option plotinfo.plotymargin

        self.yadjust = 0.0

        # Each new line is in z-order below the previous one. change it False

        # to have lines paint above the previous line

        self.zdown = True

        # Rotation of the date labes on the x axis

        self.tickrotation = 15

        # How many "subparts" takes a major chart (datas) in the overall chart

        # This is proportional to the total number of subcharts

        self.rowsmajor = 5

        # How many "subparts" takes a minor chart (indicators/observers) in the

        # overall chart. This is proportional to the total number of subcharts

        # Together with rowsmajor, this defines a proportion ratio betwen data

        # charts and indicators/observers charts

        self.rowsminor = 1

        # Distance in between subcharts

        self.plotdist = 0.0

        # Have a grid in the background of all charts

        self.grid = True

        # Default plotstyle for the OHLC bars which (line -> line on close)

        # Other options: 'bar' and 'candle'

        self.style = 'line'

        # Default color for the 'line on close' plot

        self.loc = 'black'

        # Default color for a bullish bar/candle (0.75 -> intensity of gray)

        self.barup = '0.75'

        # Default color for a bearish bar/candle

        self.bardown = 'red'

        # Level of transparency to apply to bars/cancles (NOT USED)

        self.bartrans = 1.0

        # Wether the candlesticks have to be filled or be transparent

        self.barupfill = True

        self.bardownfill = True

        # Wether the candlesticks have to be filled or be transparent

        self.fillalpha = 0.20

        # Wether to plot volume or not. Note:if the data in question has no

        # volume values, volume plotting will be skipped even if this is True

        self.volume = True

        # Wether to overlay the volume on the data or use a separate subchart

        self.voloverlay = True

        # Scaling of the volume to the data when plotting as overlay

        self.volscaling = 0.33

        # Pushing overlay volume up for better visibiliy. Experimentation

        # needed if the volume and data overlap too much

        self.volpushup = 0.00

        # Default colour for the volume of a bullish day

        self.volup = '#aaaaaa'# 0.66 of gray

        # Default colour for the volume of a bearish day

        self.voldown = '#cc6073'# (204, 96, 115)

        # Transparency to apply to the volume when overlaying

        self.voltrans = 0.50

        # Transparency for text labels (NOT USED CURRENTLY)

        self.subtxttrans = 0.66

        # Default font text size for labels on the chart

        self.subtxtsize = 9

        # Transparency for the legend (NOT USED CURRENTLY)

        self.legendtrans = 0.25

        # Wether indicators have a leged displaey in their charts

        self.legendind = True

        # Location of the legend for indicators (see matplotlib)

        self.legendindloc = 'upper left'

        # Plot the last value of a line after the Object name

        self.linevalues = True

        # Plot a tag at the end of each line with the last value

        self.valuetags = True

        # Default color for horizontal lines (see plotinfo.plothlines)

        self.hlinescolor = '0.66'# shade of gray

        # Default style for horizontal lines

        self.hlinesstyle = '--'

        # Default width for horizontal lines

        self.hlineswidth = 1.0

        # Default color scheme: Tableau 10

        self.lcolors = tableau10

        # strftime Format string for the display of ticks on the x axis

        self.fmt_x_ticks = None

        # strftime Format string for the display of data points values

        self.fmt_x_data = None
